package com.example.gradu.graduationproject.controller;


import com.gearwenxin.client.PromptBotClient;
import com.gearwenxin.client.ernie.ErnieBot4Client;
import com.gearwenxin.client.stable.StableDiffusionXLClient;
import com.gearwenxin.entity.enums.SamplerType;
import com.gearwenxin.entity.request.ImageBaseRequest;
import com.gearwenxin.entity.response.ChatResponse;
import com.gearwenxin.entity.response.ImageResponse;
import com.gearwenxin.entity.response.PromptResponse;
import org.springframework.http.MediaType;
import org.springframework.web.bind.annotation.*;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;
import reactor.core.publisher.Mono;
import com.gearwenxin.entity.chatmodel.ChatPromptRequest;
import jakarta.annotation.Resource;


import java.util.HashMap;
import java.util.Map;


/**
 * 百度文心一言调用接口(已废弃）
 */

@RestController
@RequestMapping("chatAi")
public class ChatController {

    // 要调用的模型的客户端（示例为文心4.0）
    @Resource
    private ErnieBot4Client ernieBot4Client;
    @Resource
    private StableDiffusionXLClient stableDiffusionXLClient;
    @Resource
    private PromptBotClient promptClient;

    // 单次对话
    @PostMapping("/chat")
    public Mono<ChatResponse> chatSingle(@RequestParam  String msg) {
        return ernieBot4Client.chatSingle(msg);
    }

    // 流式返回，连续对话（SSE形式）
    @GetMapping(value = "/stream/chats", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public SseEmitter chatSingleSSE(@RequestParam String msg) {
        String chatUID = "test-user-1001";
        SseEmitter emitter = new SseEmitter();
        ernieBot4Client.chatContOfStream(msg, chatUID).subscribe(response -> {
            try {
                emitter.send(SseEmitter.event().data(response.getResult()));
            } catch (Exception e) {
                emitter.completeWithError(e);
            }
        }, emitter::completeWithError, emitter::complete);
        return emitter;
    }

    // Prompt模板
    @PostMapping("/prompt")
    public Mono<PromptResponse> chatSingle(int id) {
        Map<String, String> map = new HashMap<>();
        map.put("article", "人们都说：“桂林山水甲天下。”我们乘着木船，荡漾在漓江上，来观赏桂林的山水。\n" +
                "我看见过波澜壮阔的大海，玩赏过水平如镜的西湖，却从没看见过漓江这样的水。漓江的水真静啊，静得让你感觉不到它在流动；漓江的水真清啊，清得可以看见江底的沙石；漓江的水真绿啊，绿得仿佛那是一块无瑕的翡翠。船桨激起的微波扩散出一道道水纹，才让你感觉到船在前进，岸在后移。\n" +
                "我攀登过峰峦雄伟的泰山，游览过红叶似火的香山，却从没看见过桂林这一带的.山，桂林的山真奇啊，一座座拔地而起，各不相连，像老人，像巨象，像骆驼，奇峰罗列，形态万千；桂林的山真秀啊，像翠绿的屏障，像新生的竹笋，色彩明丽，倒映水中；桂林的山真险啊，危峰兀立，怪石嶙峋，好像一不小心就会栽倒下来。\n" +
                "这样的山围绕着这样的水，这样的水倒映着这样的山，再加上空中云雾迷蒙，山间绿树红花，江上竹筏小舟，让你感到像是走进了连绵不断的画卷，真是“舟行碧波上，人在画中游”。");
        map.put("number", "20");
        ChatPromptRequest promptRequest = new ChatPromptRequest();
        promptRequest.setId(id);
        promptRequest.setParamMap(map);

        return promptClient.chatPrompt(promptRequest);
    }

    // 文生图
    @PostMapping("/image")
    public Mono<ImageResponse> chatImage() {
        ImageBaseRequest imageBaseRequest = ImageBaseRequest.builder()
                // 提示词
                .prompt("一个头发中分背带裤的人")
                // 大小
                .size("1024x1024")
                // 反向提示词（不包含什么）
                .negativePrompt("鸡")
                // 生成图片数量（1-4）
                .n(1)
                // 迭代轮次（10-50）
                .steps(20)
                // 采样方式
                .samplerIndex(SamplerType.Euler_A.getValue())
                .userId("1001")
                .build();

        return stableDiffusionXLClient.chatImage(imageBaseRequest);
    }

}